import json

from sdv.metadata import MultiTableMetadata
from sqlalchemy import text

from tool_utils import get_sql_lite_tool
from tool_utils import get_database_metadata
from tool_utils import get_single_filed_type_dict, get_all_table_and_relationships, update_database_metadata


def database_table_relationship_mange(st):
    st.title("表关系管理")

    if 'new_parent' not in st.session_state:
        st.session_state.new_parent = None
    if 'new_parent_id' not in st.session_state:
        st.session_state.new_parent_id = ''

    if 'new_child' not in st.session_state:
        st.session_state.new_parent = None
    if 'new_child_foreign_key' not in st.session_state:
        st.session_state.new_parent_id = None

    # 获取所有数据源
    sql_lite_tool = get_sql_lite_tool()
    with sql_lite_tool.engine.connect() as conn:
        # 获取所有数据源
        data_sources = conn.execute(text('''SELECT id, name 
            FROM database_connections dc
            WHERE EXISTS (
                SELECT 1 
                FROM database_all_table_metadata datm 
                WHERE datm.database_id = dc.id  
            );''')).fetchall()

    ## 需要过滤掉没有
    if not data_sources:
        st.warning("没有找到任何数据源，请先添加数据源")
        return

    # 数据源选择
    selected_data_source = st.selectbox("选择数据源",
                                        options=[ds[1] for ds in data_sources],
                                        format_func=lambda x: x)

    if selected_data_source:
        # 获取选中数据源的id
        with sql_lite_tool.engine.connect() as conn:
            data_source_id = \
                conn.execute(
                    text(f"SELECT id FROM database_connections WHERE name = '{selected_data_source}'")).fetchone()[
                    0]

            ## 获取元数据

            all_table_metadata = get_database_metadata(data_source_id)
            relationships = []
            if all_table_metadata:

                my_metadata = get_all_table_and_relationships(str(all_table_metadata))
                table_name_list = my_metadata['table_name_list']
                relationships = my_metadata['relationships']
                relationships = sorted(
                    relationships,
                    key=lambda x: (x["child_table_name"], x["parent_table_name"])
                )
                table_name_primary_key_dict = my_metadata['table_name_primary_key_dict']
                table_name_columns_dict = my_metadata['table_name_columns_dict']

            else:
                st.rerun()

        # 显示表关系
        st.subheader("表关系列表")

        for index, relationship in enumerate(relationships):
            col1, col2, col3, col4, col5, col6 = st.columns([2, 2, 2, 2, 1, 1])
            parent_table_name_old = relationship['parent_table_name']
            child_table_name_old = relationship['child_table_name']
            parent_id_old = relationship['parent_primary_key']
            child_id_old = relationship['child_foreign_key']
            with col1:

                st.text_input("子表", value=child_table_name_old,
                              key=f"child_{index}",
                              disabled=True)

            columns = table_name_columns_dict.get(child_table_name_old)
            with col2:
                child_id = st.selectbox("子表外键",
                                        options=columns,
                                        index=columns.index(child_id_old),
                                        key=f"child_id_{index}")

            with col3:
                parent_table_name = st.selectbox("父表",
                                                 options=[t for t in table_name_list],
                                                 index=table_name_list.index(parent_table_name_old),
                                                 key=f"parent_{index}")

            with col4:
                parent_id = st.text_input("父表主键", value=table_name_primary_key_dict.get(parent_table_name),
                                          key=f"parent_id_{index}",
                                          disabled=True)

            with col5:
                st.write('')
                st.write('')


                if st.button("🗑️ 删除", key=f"delete_{index}"):
                    all_table_metadata.remove_relationship(parent_table_name_old, child_table_name_old)
                    update_database_metadata(database_id=data_source_id, metadata=all_table_metadata)
                    st.success(f"已删除关系: {child_table_name_old} -> {parent_table_name_old}")
                    st.rerun()

            with col6:
                st.write('')
                st.write('')
                if st.button("✔ 保存", key=f"save_{index}"):
                    all_table_metadata.remove_relationship(parent_table_name_old, child_table_name_old)
                    all_table_metadata.add_relationship(parent_table_name, child_table_name_old, parent_id, child_id)
                    update_database_metadata(database_id=data_source_id, metadata=all_table_metadata)
                    # 清空所有值
                    st.session_state.new_child = None
                    st.session_state.new_child_foreign_key = None
                    st.session_state.new_parent = None
                    st.session_state.new_parent_id = ''
                    st.rerun()

        # 添加新关系
        st.subheader("添加新关系")
        col1, col2, col3, col4, col5 = st.columns([1, 1, 1, 1, 1])

        with col1:
            new_child = st.selectbox("选择子表",
                                     options=table_name_list,
                                     key="new_child",
                                     index=None, )

        child_columns = table_name_columns_dict.get(new_child)
        with col2:
            new_child_foreign_key = st.selectbox("选择子表外键",
                                                 options=child_columns,
                                                 key="new_child_foreign_key", index=None, )
        with col3:
            new_parent = st.selectbox("选择父表",
                                      options=table_name_list,
                                      key="new_parent", index=None, on_change=lambda: st.session_state.update({
                    "new_parent": st.session_state.new_parent,
                    "new_parent_id": table_name_primary_key_dict.get(st.session_state.new_parent)
                }))

        with col4:

            new_parent_id = st.text_input("父表主键", value=st.session_state.new_parent_id,
                                          key="new_parent_id",
                                          disabled=True)

        with col5:
            st.write('')
            st.write('')
            if st.button("✔ 确认添加"):
                try:
                    all_table_metadata.add_relationship(new_parent, new_child, new_parent_id, new_child_foreign_key)
                    update_database_metadata(database_id=data_source_id, metadata=all_table_metadata)

                    st.success(f"已添加新关系: {new_child} -> {new_parent}")
                    st.rerun()
                except Exception as e:
                    st.error(f"添加关系失败: {str(e)}")

        st.subheader("关系展示")
        if st.button("展示结构", key=f"show_relationship_{data_source_id}"):
            st.write(all_table_metadata.visualize(
                show_table_details='full',
                show_relationship_labels=True,
            ))


def newnew(st):
    st.subheader("Test")
